from sklearn.metrics import confusion_matrix
import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE


class ImgPainter:
    @staticmethod
    def img_line(y, x=None, path=None, title=None, max_index=None):
        if x is None:
            x = np.arange(len(y))

            fig, ax = plt.subplots()
            ax.plot(x, y)
            if max_index is not None:
                ax.scatter(max_index, y[max_index], color='red', zorder=5)
            ax.grid()
            if title is not None:
                ax.set_title(title)
            if path is not None:
                plt.savefig(path)

            plt.show()

    @staticmethod
    def tsne_plot(x, y,title):
        tsne = TSNE(n_components=2, random_state=42)
        x_tsne = tsne.fit_transform(x)

        # Plot the t-SNE visualization
        # plt.figure(figsize=(8, 6))
        colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k', 'orange', 'purple', 'brown']
        for i in np.unique(y):
            plt.scatter(x_tsne[y == i, 0], x_tsne[y == i, 1], color=colors[i % 10], label=str(i), s=1)
        plt.title(title)
        plt.legend()
        plt.show()

    @staticmethod
    def confusion_matrix_plot(y_true, y_pre):
        # 如果没有指定标签，则从 y_true 和 prey_pred 中确定
        labels = np.unique(np.concatenate((y_true, y_pre)))

        # 创建混淆矩阵
        cm = confusion_matrix(y_true, y_pre, labels=labels)

        # 绘制混淆矩阵
        plt.figure(figsize=(10, 6))
        sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=labels, yticklabels=labels)
        plt.title('Confusion Matrix')
        plt.xlabel('Predicted Class')
        plt.ylabel('True Class')
        plt.show()
